package com.jxpanda.spring.module.auth.config.properties;

import com.jxpanda.infrastructure.core.tollkit.json.JsonKit;
import com.nimbusds.jose.Algorithm;
import com.nimbusds.jose.Requirement;
import com.nimbusds.jose.jwk.*;
import com.nimbusds.jose.jwk.source.ImmutableJWKSet;
import com.nimbusds.jose.jwk.source.JWKSource;
import com.nimbusds.jose.proc.SecurityContext;
import lombok.Getter;
import lombok.RequiredArgsConstructor;
import lombok.Setter;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.boot.context.properties.ConfigurationProperties;
import org.springframework.core.io.ClassPathResource;
import org.springframework.security.oauth2.jose.jws.JwsAlgorithm;
import org.springframework.security.oauth2.jose.jws.MacAlgorithm;
import org.springframework.security.oauth2.jose.jws.SignatureAlgorithm;

import java.io.File;
import java.io.IOException;
import java.net.URL;
import java.security.*;
import java.security.interfaces.ECPublicKey;
import java.security.interfaces.RSAPublicKey;
import java.security.spec.ECGenParameterSpec;
import java.text.ParseException;
import java.util.Arrays;
import java.util.List;
import java.util.Objects;
import java.util.UUID;
import java.util.function.Function;

@Getter
@Setter
@Slf4j
@SuppressWarnings("HttpUrlsUsage")
@ConfigurationProperties(prefix = "panda.spring.auth.token.jwt")
public class JwtProperties {

    private static final List<String> DEFAULT_AUDIENCE = List.of("local");

    /**
     * Jwt算法，默认值是：ES384
     */
    private JwtAlgorithm algorithm = JwtAlgorithm.ES384;

    /**
     * Jwt issuer，默认值是：取spring.application.name的值，如果为空则取localhost
     */
    @Value("${panda.spring.auth.token.jwt.issuer:${spring.application.name:localhost}}")
    private String issuer;

    /**
     * Jwt audience，默认值是：local
     */
    private List<String> audience = DEFAULT_AUDIENCE;

    /**
     * 加载Jwt的资源路径，可以是一个URL 或一个文件路径。
     * 使用文件路径的话，必须是一个jwks的json文件。
     * 使用URL的话必须确定URL是有效的
     * <p>
     * 如果为空，则自动指定的算法生成一个 JWKSource。
     */
    private String jwkSetUri;

    /**
     * 是否启用
     */
    private boolean enable = false;

    public JWKSource<SecurityContext> loadJwkSource() {
        if (Objects.isNull(this.jwkSetUri)) {
            // jwkSetUri 为空时，基于算法生成 JWKSource
            JWKSource<SecurityContext> jwkSource = this.getAlgorithm().generateJwkSource();
            if (log.isDebugEnabled()) {
                ImmutableJWKSet<SecurityContext> jwkSet = (ImmutableJWKSet<SecurityContext>) jwkSource;
                log.debug("[PANDA SPRING AUTH]Jwt密钥 已生成: {}", JsonKit.toJson(jwkSet.getJWKSet().toJSONObject(false)));
            }
            return jwkSource;
        }

        try {
            JWKSet jwkSet = UriProtocol.load(jwkSetUri);
            if (!jwkSet.isEmpty()) {
                this.algorithm = JwtAlgorithm.valueOf(jwkSet.getKeys().get(0).getAlgorithm().getName());
            }
            return new ImmutableJWKSet<>(jwkSet);
        } catch (Exception e) {
            throw new RuntimeException("加载 JWKSet 失败: " + e.getMessage(), e);
        }
    }

    @RequiredArgsConstructor
    private enum UriProtocol {
        HTTP("http://", false, UriProtocol::loadFromUrl),
        HTTPS("https://", false, UriProtocol::loadFromUrl),

        FILE("file://", true, UriProtocol::loadFromFile),

        CLASSPATH("classpath:", true, UriProtocol::loadFromResource),
        CLASSPATH_ALL("classpath*:", true, UriProtocol::loadFromResource);


        /**
         * 前缀集合
         */
        private final String prefix;

        /**
         * 加载的时候是否需要清空协议头
         */
        private final boolean clearPrefix;

        /**
         * 加载函数
         */
        private final Function<String, JWKSet> loader;


        private static JWKSet loadFromUrl(String url) {
            try {
                return JWKSet.load(new URL(url));
            } catch (IOException | ParseException e) {
                throw new RuntimeException(e);
            }
        }

        private static JWKSet loadFromFile(String filePath) {
            try {
                return JWKSet.load(new File(filePath));
            } catch (IOException | ParseException e) {
                throw new RuntimeException(e);
            }
        }

        private static JWKSet loadFromResource(String resource) {
            try {
                return JWKSet.load(new ClassPathResource(resource).getInputStream());
            } catch (IOException | ParseException e) {
                throw new RuntimeException(e);
            }
        }

        public static JWKSet load(String uri) {
            return Arrays.stream(UriProtocol.values())
                    .filter(uriProtocol -> uri.startsWith(uriProtocol.prefix))
                    .findFirst()
                    .map(uriProtocol -> {
                        String uriString = uriProtocol.clearPrefix ? uri.replace(uriProtocol.prefix, "") : uri;
                        return uriProtocol.loader.apply(uriString);
                    })
                    .orElseGet(() -> loadFromFile(uri));
        }

    }

    public enum JwtAlgorithmFamily {
        HMAC_SHA, RSA, EC, ED;
    }

    @Getter
    @RequiredArgsConstructor
    public enum JwtAlgorithm {
        HS256(JwtAlgorithmFamily.HMAC_SHA, MacAlgorithm.HS256, 32),
        HS384(JwtAlgorithmFamily.HMAC_SHA, MacAlgorithm.HS384, 48),
        HS512(JwtAlgorithmFamily.HMAC_SHA, MacAlgorithm.HS512, 64),
        RS256(JwtAlgorithmFamily.RSA, SignatureAlgorithm.RS256, 2048),
        RS384(JwtAlgorithmFamily.RSA, SignatureAlgorithm.RS384, 3072),
        RS512(JwtAlgorithmFamily.RSA, SignatureAlgorithm.RS512, 4096),
        ES256(JwtAlgorithmFamily.EC, SignatureAlgorithm.ES256, "secp256r1"),
        ES384(JwtAlgorithmFamily.EC, SignatureAlgorithm.ES384, "secp384r1"),
        ES512(JwtAlgorithmFamily.EC, SignatureAlgorithm.ES512, "secp521r1"),
        PS256(JwtAlgorithmFamily.RSA, SignatureAlgorithm.PS256, 2048),
        PS384(JwtAlgorithmFamily.RSA, SignatureAlgorithm.PS384, 3072),
        PS512(JwtAlgorithmFamily.RSA, SignatureAlgorithm.PS512, 4096);
        // 以下3个算法暂时不支持
//        EdDSA(ED, JWSAlgorithm.EdDSA, ""),
//        Ed25519(ED, JWSAlgorithm.Ed25519, ""),
//        Ed448(ED, JWSAlgorithm.Ed448, "");


        private final JwtAlgorithmFamily family;

        private final JwsAlgorithm jwsAlgorithm;

        private final Object keySpec;

        private JWKSource<SecurityContext> generateJwkSource() {
            return switch (this.family) {
                case RSA -> {
                    try {
                        yield generateRsaJwkSource();
                    } catch (NoSuchAlgorithmException e) {
                        throw new RuntimeException(e);
                    }
                }
                case EC -> {
                    try {
                        yield generateEcJwkSource();
                    } catch (NoSuchAlgorithmException | InvalidAlgorithmParameterException e) {
                        throw new RuntimeException(e);
                    }
                }
                case HMAC_SHA -> generateHmacJwkSource();
                default -> throw new IllegalArgumentException("Unsupported algorithm family: " + this.family);
            };
        }


        /**
         * 生成 RSA 类型的 JWKSource
         */
        private JWKSource<SecurityContext> generateRsaJwkSource() throws NoSuchAlgorithmException {
            // 生成 RSA 密钥对
            KeyPairGenerator keyPairGenerator = KeyPairGenerator.getInstance("RSA");
            // 可根据需要调整密钥大小
            keyPairGenerator.initialize((int) this.keySpec);
            KeyPair keyPair = keyPairGenerator.generateKeyPair();
            // 创建 RSAKey
            RSAKey rsaKey = new RSAKey.Builder((RSAPublicKey) keyPair.getPublic())
                    .privateKey(keyPair.getPrivate())
                    .algorithm(new Algorithm(this.getJwsAlgorithm().getName(), Requirement.REQUIRED))
                    // 生成唯一的 key ID
                    .keyID(UUID.randomUUID().toString())
                    .build();
            // 创建 JWKSet
            JWKSet jwkSet = new JWKSet(rsaKey);
            return new ImmutableJWKSet<>(jwkSet);
        }

        /**
         * 生成 EC 类型的 JWKSource
         */
        private JWKSource<SecurityContext> generateEcJwkSource() throws NoSuchAlgorithmException, InvalidAlgorithmParameterException {
            // 生成 EC 密钥对
            KeyPairGenerator keyPairGenerator = KeyPairGenerator.getInstance("EC");
            ECGenParameterSpec ecSpec = new ECGenParameterSpec((String) this.keySpec);
            keyPairGenerator.initialize(ecSpec);
            KeyPair keyPair = keyPairGenerator.generateKeyPair();

            // 创建 ECKey
            ECKey ecKey = new ECKey.Builder(Curve.forECParameterSpec(((ECPublicKey) keyPair.getPublic()).getParams()), (ECPublicKey) keyPair.getPublic())
                    .privateKey(keyPair.getPrivate())
                    .algorithm(new Algorithm(this.getJwsAlgorithm().getName(), Requirement.REQUIRED))
                    // 生成唯一的 key ID
                    .keyID(UUID.randomUUID().toString())
                    .build();

            // 创建 JWKSet
            JWKSet jwkSet = new JWKSet(ecKey);
            return new ImmutableJWKSet<>(jwkSet);
        }

        /**
         * 生成 HMAC 类型的 JWKSource
         */
        private JWKSource<SecurityContext> generateHmacJwkSource() {

            // 生成随机密钥
            byte[] secret = new byte[(int) this.getKeySpec()];
            SecureRandom secureRandom = new SecureRandom();
            secureRandom.nextBytes(secret);

            // 创建 OctetSequenceKey
            OctetSequenceKey hmacKey = new OctetSequenceKey.Builder(secret)
                    .algorithm(new Algorithm(this.getJwsAlgorithm().getName(), Requirement.REQUIRED))
                    // 生成唯一的 key ID
                    .keyID(UUID.randomUUID().toString())
                    .build();

            // 创建 JWKSet
            JWKSet jwkSet = new JWKSet(hmacKey);

            return new ImmutableJWKSet<>(jwkSet);
        }
    }

}
